__author__ = "kiana";

import jieba;
import jieba.analyse;
from collections import Counter;

'''
分词、词频统计测试
'''

# 目标文件
target_file = "E:/temp/冒牌大英雄.txt";
# 停用词表文件
stop_words_file = "../resources/jieba_stop_words.txt";
# 自定义分词字典文件
user_dict_file = "../resources/jieba_user_dict.txt";

# 针对所有内容的关键字词频统计
def get_keyword_frequency(keyword_list=[]):
    print("==== 获取目标字符串 ====");
    content = "";
    with open(file=target_file, mode="r", encoding="utf-8") as f:
        content = f.read().strip().replace("\n", "").replace(" ", "");
    print(content[0:20] + "......");

    print("==== 获取停用词表 ====");
    stop_word_list = [];
    with open(file=stop_words_file, mode="r", encoding="utf-8") as f:
        # 每行去空格
        stop_word_list = list(line.strip() for line in f.readlines());
    print(stop_word_list);

    print("=== 获取自定义分词词典 ====");
    '''
    添加 jieba 里未能识别的词语
    一个词占一行，每一行分成三部分，一部分为词语，一部分为词频(可省略)，最后为词性(可以省略)
    '''
    jieba.load_userdict(user_dict_file);

    print("==== 分词 ====");
    # 分词，cut 返回 generator，lcut 返回 list
    cut_word_list = jieba.lcut(content); # 精确模式
    # cut_word_list = jieba.lcut(content, cut_all=True); # 全模式
    # cut_word_list = jieba.lcut_for_search(content); # 搜索引擎模式
    print(cut_word_list);

    print("==== 词频统计 ====");
    # 方式一
    word_frequency_map = {};
    for item in cut_word_list:
        if item not in stop_word_list:
            word_frequency_map[item] = word_frequency_map.get(item, 0) + 1;
    # 将 map 转化为 list，并按词频大小排序
    word_frequency_list = list(word_frequency_map.items());
    word_frequency_list.sort(key=lambda x:x[1], reverse=True);
    print(word_frequency_list);
    # 方式二，返回 [('，', 5932), ('的', 3177),...] 格式
    counter = Counter(cut_word_list);
    print(counter.most_common());

    print("==== 提取关键字 ====");
    if len(keyword_list) == 0:
        jieba.analyse.set_stop_words(stop_words_file);
        # 基于 TF/IDF 算法，返回权重最大的关键词，默认值为 100，withWeight 是否返回权重值
        keyword_list = jieba.analyse.extract_tags(content, topK=100, withWeight=False);
        print(keyword_list);
        # 基于 TextRank 算法，返回权重最大的关键词，默认值为 100，withWeight 是否返回权重值
        # keyword_list = jieba.analyse.textrank(content, topK=100, withWeight=False);
        # print(keyword_list);

    print("==== 提取含有关键字的词频统计数据 ==== ");
    keyword_frequency_list = list(item for item in word_frequency_list if item[0] in keyword_list);
    print(keyword_frequency_list);

    return word_frequency_list, keyword_frequency_list;

if __name__ == "__main__":
    get_keyword_frequency();